#!/usr/bin/env bash

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0

# A utility script to download component wheels from GitHub Actions artifacts.
# This script reuses the same logic that was in release.yml to maintain consistency.

set -euo pipefail

# Check required arguments
if [[ $# -lt 3 ]]; then
    echo "Usage: $0 <run-id> <component> <repository> [output-dir]" >&2
    echo "  run-id: The GitHub Actions run ID containing the artifacts" >&2
    echo "  component: The component name pattern to download (e.g., cuda-core, cuda-bindings)" >&2
    echo "  repository: The GitHub repository (e.g., NVIDIA/cuda-python)" >&2
    echo "  output-dir: Optional output directory (default: ./dist)" >&2
    exit 1
fi

RUN_ID="$1"
COMPONENT="$2"
REPOSITORY="$3"
OUTPUT_DIR="${4:-./dist}"

# Ensure we have a GitHub token
if [[ -z "${GH_TOKEN:-}" ]]; then
    echo "Error: GH_TOKEN environment variable is required"
    exit 1
fi

echo "Downloading wheels for component: $COMPONENT from run: $RUN_ID"

# Download component wheels using the same logic as release.yml
if [[ "$COMPONENT" == "all" ]]; then
    # Download all component patterns
    gh run download "$RUN_ID" -p "cuda-*" -R "$REPOSITORY"
else
    gh run download "$RUN_ID" -p "${COMPONENT}*" -R "$REPOSITORY"
fi

# Create output directory
mkdir -p "$OUTPUT_DIR"

# Process downloaded artifacts
for p in cuda-*
do
    if [[ ! -d "$p" ]]; then
        continue
    fi

    # exclude cython test artifacts
    if [[ "${p}" == *-tests ]]; then
        echo "Skipping test artifact: $p"
        continue
    fi

    # If we're not downloading "all", only process matching component
    if [[ "$COMPONENT" != "all" && "$p" != ${COMPONENT}* ]]; then
        continue
    fi

    echo "Processing artifact: $p"
    # Move wheel files to output directory
    if [[ -d "$p" ]]; then
        find "$p" -name "*.whl" -exec mv {} "$OUTPUT_DIR/" \;
    fi
done

# Clean up artifact directories
rm -rf cuda-*

echo "Downloaded wheels to: $OUTPUT_DIR"
ls -la "$OUTPUT_DIR"
